package org.example;

import dev.langchain4j.agent.tool.P;
import dev.langchain4j.agent.tool.ToolExecutionRequest;
import dev.langchain4j.agent.tool.ToolSpecification;
import dev.langchain4j.agent.tool.ToolSpecifications;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.output.Response;

import java.time.LocalDate;
import java.time.ZoneOffset;
import java.time.format.DateTimeFormatter;
import java.util.ArrayList;
import java.util.List;

/**
 * @author fish_temp_author
 * @since fish_temp_since
 */
public class ToolDemo {
    public static void main(String[] args) throws NoSuchMethodException {

        ChatLanguageModel openAiChatModel = OpenAiChatModel.builder()
                .apiKey("demo")
                .modelName("gpt-4o-mini")
                .build();

        // 传入方法，并询问今天的日期是多少
        ArrayList<ToolSpecification> toolSpecifications = new ArrayList<>();
        toolSpecifications.add(ToolSpecifications.toolSpecificationFrom(Tool.class.getMethod("getCurrentDate")));
        toolSpecifications.add(ToolSpecifications.toolSpecificationFrom(Tool.class.getMethod("getUserName")));
        toolSpecifications.add(ToolSpecifications.toolSpecificationFrom(Tool.class.getMethod("getBmi", String.class, String.class)));

        Response<AiMessage> aiMessageResponse = openAiChatModel.generate(List.of(UserMessage.from("今天是几号")), toolSpecifications);
        if (aiMessageResponse.content()
                .hasToolExecutionRequests()) {
            List<ToolExecutionRequest> toolExecutionRequests = aiMessageResponse.content().toolExecutionRequests();
            for (ToolExecutionRequest toolExecutionRequest : toolExecutionRequests) {
                String toolName = toolExecutionRequest.name();
                String toolInput = toolExecutionRequest.arguments();
                System.out.println(toolName);
                System.out.println(toolInput);
            }
        }

        Response<AiMessage> aiMessageResponse2 = openAiChatModel.generate(List.of(UserMessage.from("今天是几号")), toolSpecifications);
        if (aiMessageResponse2.content()
                .hasToolExecutionRequests()) {
            List<ToolExecutionRequest> toolExecutionRequests = aiMessageResponse2.content().toolExecutionRequests();
            for (ToolExecutionRequest toolExecutionRequest : toolExecutionRequests) {
                String toolName = toolExecutionRequest.name();
                String toolInput = toolExecutionRequest.arguments();
                System.out.println(toolName);
                System.out.println(toolInput);
            }
        }
        Response<AiMessage> aiMessageResponse3 = openAiChatModel.generate(List.of(UserMessage.from("你能帮我计算bmi吗，身高180cm，体重100KG")), toolSpecifications);
        if (aiMessageResponse3.content().hasToolExecutionRequests()) {
            List<ToolExecutionRequest> toolExecutionRequests = aiMessageResponse3.content().toolExecutionRequests();
            for (ToolExecutionRequest toolExecutionRequest : toolExecutionRequests) {
                String toolName = toolExecutionRequest.name();
                String toolInput = toolExecutionRequest.arguments();
                System.out.println(toolName);
                System.out.println(toolInput);
            }
        } else {
            System.out.println(aiMessageResponse3.content().text());
        }
    }


    public static class Tool {
        @dev.langchain4j.agent.tool.Tool(name = "get_current_date", value = "获取当前的日期")
        public static String getCurrentDate() {
            return LocalDate.now(ZoneOffset.ofHours(8)).format(DateTimeFormatter.ofPattern("yyyy-MM-dd"));
        }

        @dev.langchain4j.agent.tool.Tool(name = "get_username", value = "获取名字")
        public static String getUserName() {
            return "我叫谢阿花";
        }

        // 根据身高体重计算bmi
        @dev.langchain4j.agent.tool.Tool(name = "get_bmi", value = "根据身高体重计算bmi")
        public static String getBmi(@P("体重（千克）") String height, @P("身高（厘米）") String weight) {
            double bmi = Double.parseDouble(weight) / (Double.parseDouble(height) * Double.parseDouble(height));
            return String.valueOf(bmi);
        }
    }
}
